#!/bin/bash
set -x

# 接受参数，控制是否启用hierarchical cache（1启用，0或无参数不启用）
param=$1
# 参数处理：未传参或非1则默认不启用
if [ -z "$param" ] || [ "$param" -ne 1 ]; then
    param=0
fi

# 修正端口一致：服务器启动在9000端口，测试连接也应使用9000
deepseek_url=http://127.0.0.1:9000
model=/data2/public-models/Qwen3-235B-A22B-FP8
tokenizer=/data2/public-models/Qwen3-235B-A22B-FP8
dataset_path=/data1/dataset/ShareGPT_V3_unfiltered_cleaned_split.json
backend=sglang  # sglang/sglang-oai

##########################
# 测试参数配置
ISLS="20480"
OSLS="1024"
CON="80 100 150 170"
# 每次启动服务器后跑3遍测试
RUN_TIMES=3
##########################

# 启动服务器函数
start_server() {
    echo "Starting sglang server with param=$param..."
    if [ "$param" -eq 1 ]; then
        # 参数为1时，启用hierarchical cache
        python3 -m sglang.launch_server \
            --model-path $model \
            --tensor-parallel-size 4 \
            --host 0.0.0.0 \
            --port 9000 \
            --enable-hierarchical-cache \
            --hicache-size 100 \
            --mem-fraction-static 0.9 &
    else
        # 参数为0或未传参时，不启用hierarchical cache
        python3 -m sglang.launch_server \
            --model-path $model \
            --tensor-parallel-size 4 \
            --host 0.0.0.0 \
            --port 9000 \
            --mem-fraction-static 0.9 &
    fi
    SERVER_PID=$!
    echo "Server started with PID: $SERVER_PID"
    
    # 等待服务器启动完成（根据模型大小调整等待时间，建议至少60s）
    echo "Waiting for server to be ready..."
    sleep 120s
}

# 关闭服务器函数
stop_server() {
    if [ -n "$SERVER_PID" ]; then
        echo "Stopping server with PID: $SERVER_PID"
        kill $SERVER_PID
        # 等待进程退出
        wait $SERVER_PID 2>/dev/null
        echo "Server stopped"
        SERVER_PID=""
    fi
}

# 执行测试函数（参数：当前运行次数）
run_tests() {
    local run_num=$1
    echo "Starting test run $run_num..."
    
    for ISL in $ISLS; do
        for OSL in $OSLS; do
            SYS_TIME=$(date "+%Y%m%d_%H%M%S")
            # 日志名包含运行次数，便于区分
            LOG="temp_run${run_num}_${SYS_TIME}"
            echo "Log file: ${LOG}.log"
            
            for con in $CON; do
                num_prompt=200
                # 确保num_prompt不超过200（原逻辑保留）
                if [ "$num_prompt" -gt 200 ]; then
                    num_prompt=200
                fi
                
                echo "Running test - ISL: $ISL, OSL: $OSL, CON: $con, Run: $run_num"
                python3 -m sglang.bench_serving \
                    --backend $backend \
                    --base-url $deepseek_url \
                    --model $model \
                    --tokenizer=$tokenizer \
                    --dataset-name random \
                    --random-input-len $ISL \
                    --random-output-len $OSL \
                    --num-prompt $num_prompt \
                    --max-concurrency $con \
                    --dataset-path $dataset_path \
                    2>&1 | tee -a ${LOG}.log
            done
            sleep 2s  # 测试间短暂休息
        done
    done
}

# 主流程：启动服务器 -> 跑3遍测试 -> 关闭服务器
start_server

# 执行3遍测试
for (( run=1; run<=$RUN_TIMES; run++ )); do
    run_tests $run
done

# 测试完成后关闭服务器
stop_server

echo "All tests completed"